'''
Created on Nov 5, 2011

@author: alebalbin
'''
import numpy as np
from collections import defaultdict

class tcgaRecord:
    def  __init__(self,fields,attributes):
        self.myatt=defaultdict()
        self.att=[]
        for f,a in zip(fields,attributes):
            self.myatt[f]=a
            self.att.append(f)
            
        self.submitted_sample_id=self.myatt["submitted_sample_id"]
        self.subjectid=self.submitted_sample_id.split('-')[:2]
        # Data types for this subject
        ntypes=4
        self.mydata='unknown'
        self.data_types=np.zeros(ntypes)
        self.determine_data_type()
        
        # Summary of Data types for this subjectid
        self.all_types=0
        self.exomeAndtrans=0
        self.genomeAndtrans=0
        self.ExoGenoAndtrans=0
        self.ExoGenoAndtransAndMirna=0
        self.exome=0
        self.trans=0
        
    def print_myatt(self):
        return self.att
    def print_myself(self):
        l=[]
        for f in self.att:
            l.append( self.myatt[f] )
        return l
    
    def determine_data_type(self):
        '''
          exome,
          mirna
          DNASeq_whole
          DNASeq_exome
          DNASeq_capture
          translated_to_genomic
        '''
        filename = self.myatt["filename"]
        if filename.find('exome') >= 0 or filename.find('DNASeq_capture') >=0:
            self.data_types[0] =1 #self.exome=1
            self.mydata='exome'
        elif filename.find('mirna') >=0:
            self.data_types[3] =1 #self.mirna=1
            self.mydata='mirna'
        elif filename.find('DNASeq_whole') >=0:
            self.data_types[1] =1 #self.genome=1
            self.mydata='genome'
        elif filename.find('translated_to_genomic') >=0:
            self.data_types[2] =1 #self.transcriptome =1
            self.mydata='transcriptome'
    
    def what_data_ihave(self):
        return self.mydata
    def what_disease(self):
        return self.myatt['short_disease_type']
    def is_tumor(self):
        '''
        sampleid=TCGA-02-0003-01A-01D-1490-08
        TCGA=project, 02=source_site, 0003=patient_id,01A=sample_type. if starts with 0 tumor, if starts with 1 normal. 01D=analyte, D sample type-D=DNA,
        1490=plate,08=center
        '''        
        id = self.myatt['submitted_sample_id']
        sp=id.split('-')[3]
        if int(sp[0])==0:
            return 'tumor'
        elif int(sp[0])==1:
            return 'normal'
    def what_is_mysize(self):
        return self.myatt["size_in_GB"]


def what_data_subjecthave(tcgaRecordList):
    '''
    '''
    datatypes=[]
    for r in tcgaRecordList:
        datatypes.append( r.what_data_ihave() )
    #Sort List
    datatypes.sort()
    
    '''    
    all_types=['exome','genome','transcriptome','mirna']
    exomeAndtrans=['exome','transcriptome']
    genomeAndtrans=['genome','transcriptome']
    ExoGenoAndtrans=['exome','genome','transcriptome']
    exome=['exome']
    '''
    return datatypes     

def calculate_study_size(r,bgenomes,nexomes,ntranst,nmirnas, nunknown,
                         zgenomes,zexomes,ztranst,zmirnas,znunknown,
                         texomes):
    '''
    '''
    if r.what_data_ihave() == "transcriptome":
        ntranst+=1
        ztranst+=float(r.what_is_mysize())
    if r.what_data_ihave() == "exome":
        nexomes+=1
        zexomes+=float(r.what_is_mysize())
        if r.is_tumor() =='yes':
            texomes +=1
    if r.what_data_ihave() == "genome":
        bgenomes+=1
        zgenomes += float(r.what_is_mysize())
    if r.what_data_ihave() == "mirna":
        nmirnas+=1
        zmirnas += float(r.what_is_mysize())
    if r.what_data_ihave() == "unknown":
        nunknown+=1
        znunknown += float(r.what_is_mysize())
    ol=[nexomes, ntranst, nmirnas, bgenomes, nunknown, 
        zexomes, ztranst, zmirnas, zgenomes, znunknown]
    return ol


def read_file(file):
    '''
    submitted_sample_id
    keys for the tacgRecord dictionary:
    accession    last_update    center_name    sample    filename    upload_id    file_size    upload_date    checksum    path    published    
    index    cumulative_bytes    download        submitted_sample_id    phs    version    sra_accession    short_disease_type    subject_disease_type
        subject_sex    body_site    analyte_type    histological_type    is_tumor    sorted_file_size    size_in_GB    sra_location
    
    patiendod=TCGA-02-0003-01A-01D-1490-08
    '''
    
    ifile=open(file)
    ofile=open(file+'_processed.csv','w')
    header = ifile.next().strip('\r\n').split('\t')
    subjects=defaultdict(list)
    bodysite=defaultdict(list)
    diseasetype=defaultdict(lambda: defaultdict(list))
    
    spcol=header.index("submitted_sample_id")
    bscol=header.index("short_disease_type")
    
    
    for l in ifile:
        fields=l.strip('\r\n').split('\t')
        subjectid=",".join(fields[spcol].split('-')[:3]).replace(',','-')
        subjects[subjectid].append( tcgaRecord(header,fields) )
        bodysite[fields[bscol]].append( tcgaRecord(header,fields) )
        diseasetype[fields[bscol]][subjectid].append( tcgaRecord(header,fields) )
        printheader=tcgaRecord(header,fields).print_myatt()
    
    stats_file=open('/Users/alebalbin/Documents/projects/tcga/tcga_data_inventory_withTransc.tsv','w')
    ol=['disease','Total_samples_study', 'Total_samples_included','Total_exome_sample', 'Total_Transcriptome_samples', 'Total_mirnas_samples', 'Total_genome_samples', \
        'Total_exomes_size(GB)', 'Total_transcriptome_size(GB)', 'Total_mirnas_size(GB)','Total_genome_size(GB)', 'Total_study_size(GB)']
    stats_file.write(','.join(map(str,ol)).replace(',','\t')+'\n')
    for disease, subjects in diseasetype.iteritems():
        # Stats
        total_study_size=0
        bgenomes,nexomes,ntranst,nmirnas, nunknown=0,0,0,0,0
        zgenomes,zexomes,ztranst,zmirnas,znunknown=0,0,0,0,0
        texomes=0
        nsubjects=0

        selectWithTranscriptome=False
        if selectWithTranscriptome:
            ofile=open(file+'_processed_withTransc_'+disease+'.csv','w')
        else:
            ofile=open(file+'_processed_withAll_'+disease+'.csv','w')
            
        hd = ['patiendid','dataytpes','data_available','disease','is_tumor']+printheader
        ofile.write(",".join(hd)+'\n')
        
        for sid,records in subjects.iteritems():
            datatypes = what_data_subjecthave(records)
            if selectWithTranscriptome and 'transcriptome' in datatypes:
                nsubjects+=1
                for r in records:
                    ol = [sid, ",".join(datatypes).replace(',', '_'), r.what_data_ihave(), r.what_disease(),r.is_tumor(),",".join(r.print_myself())]
                    #print ol
                    ofile.write(",".join(ol)+'\n')
                    #print sid, ",".join(datatypes).replace(',', '_'), ",".join(r.print_myself()) #.replace(',','\t')
                    # Record the file size
                    total_study_size += float(r.what_is_mysize())
                    ol1=calculate_study_size(r,bgenomes,nexomes,ntranst,nmirnas, nunknown,
                         zgenomes,zexomes,ztranst,zmirnas,znunknown,
                         texomes)
                    
            else:
                nsubjects+=1
                for r in records:
                    ol = [sid, ",".join(datatypes).replace(',', '_'), r.what_data_ihave(), r.what_disease(),r.is_tumor(),",".join(r.print_myself())]
                    #print ol
                    ofile.write(",".join(ol)+'\n')
                    #print sid, ",".join(datatypes).replace(',', '_'), ",".join(r.print_myself()) #.replace(',','\t')
                    # Record the file size
                    total_study_size += float(r.what_is_mysize())
                    ol1=calculate_study_size(r,bgenomes,nexomes,ntranst,nmirnas, nunknown,
                         zgenomes,zexomes,ztranst,zmirnas,znunknown,
                         texomes)
                
                    '''
                    if r.what_data_ihave() == "transcriptome":
                        ntranst+=1
                        ztranst+=float(r.what_is_mysize())
                    if r.what_data_ihave() == "exome":
                        nexomes+=1
                        zexomes+=float(r.what_is_mysize())
                        if r.is_tumor() =='yes':
                            texomes +=1
                    if r.what_data_ihave() == "genome":
                        bgenomes+=1
                        zgenomes += float(r.what_is_mysize())
                    if r.what_data_ihave() == "mirna":
                        nmirnas+=1
                        zmirnas += float(r.what_is_mysize())
                    if r.what_data_ihave() == "unknown":
                        nunknown+=1
                        znunknown += float(r.what_is_mysize())
                    '''

    
        ol=[disease,len(subjects), nsubjects]+ol1+[total_study_size]
        stats_file.write(','.join(map(str,ol)).replace(',','\t')+'\n')
        
        ofile.close()
    
        

file='/Users/alebalbin/Documents/projects/tcga/TCGA_phs000178_tranche_Q-1.txt'
read_file(file)

        